import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, Trainer
from datasets import load_dataset, Dataset
import numpy as np
import os
import nltk
from nltk.sentiment import SentimentIntensityAnalyzer

# 尝试下载 vader_lexicon 资源
# nltk 的 SentimentIntensityAnalyzer 需要 vader_lexicon 来进行情感分析
# 首先尝试查找该资源是否已存在，如果不存在则进行下载
try:
    nltk.data.find('sentiment/vader_lexicon.zip')
except LookupError:
    print("未找到 vader_lexicon 资源，开始下载...")
    nltk.download('vader_lexicon')
    print("vader_lexicon 资源下载完成。")
else:
    print("vader_lexicon 资源已存在。")

# 指定本地模型路径
# 这里指定了本地预训练模型的路径，后续会基于此模型进行微调
local_model_path = "C:\\Users\\zyb\\Desktop\\local_model\\bert-base-uncased"


# 从聊天记录文件中读取数据
# 此函数用于读取指定文件路径下的聊天记录数据
# 只处理包含冒号的行，将冒号后面的内容作为数据添加到列表中
def read_chat_data(file_path):
    data = []
    with open(file_path, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if line and ':' in line:  # 检查行中是否包含冒号
                data.append(line.split(':', 1)[1])
    return data


# 进行情感分析并添加到数据集中
# 使用 SentimentIntensityAnalyzer 对数据中的每条文本进行情感分析
# 根据分析结果的负向、中性、正向得分，选择得分最高的作为情感标签
def add_sentiment_analysis(data):
    sia = SentimentIntensityAnalyzer()
    sentiment_data = []
    for text in data:
        sentiment = sia.polarity_scores(text)
        sentiment_label = np.argmax([sentiment['neg'], sentiment['neu'], sentiment['pos']])
        sentiment_data.append(sentiment_label)
    return sentiment_data


# 创建数据集对象
# 将读取的聊天数据和对应的情感标签组合成一个字典
# 然后使用 Dataset.from_dict 方法将字典转换为数据集对象
def create_dataset(data, sentiment_data):
    dataset_dict = {
        'text': data,
        'sentiment': sentiment_data
    }
    return Dataset.from_dict(dataset_dict)


# 加载预训练模型和分词器
# 从指定的本地模型路径加载预训练的分词器和模型
# 设置模型的 num_labels 为 3，因为情感标签有 3 类（负向、中性、正向）
tokenizer = AutoTokenizer.from_pretrained(local_model_path)
# 设置模型的 num_labels 为 3
model = AutoModelForSequenceClassification.from_pretrained(local_model_path, num_labels=3)


# 数据预处理函数
# 对输入的数据集进行分词处理，截断和填充到最大长度
# 同时将 sentiment 列重命名为 labels
def preprocess_function(examples):
    result = tokenizer(examples["text"], truncation=True, padding='max_length')
    # 将 sentiment 列重命名为 labels
    if "sentiment" in examples:
        result["labels"] = examples["sentiment"]
        print("预处理后包含 labels 标签:", result.keys())  # 调试信息
    return result


# 自定义 Trainer 类，重写 compute_loss 方法
# 自定义的 Trainer 类继承自 Trainer，重写 compute_loss 方法
# 在计算损失时，先检查输入中是否包含标签，如果没有则抛出异常
# 然后计算交叉熵损失
class CustomTrainer(Trainer):
    def compute_loss(self, model, inputs, num_items_in_batch=None, return_outputs=False):
        print("compute_loss 输入内容:", inputs.keys())  # 调试信息
        labels = inputs.get("labels")
        # 检查 labels 是否为 None
        if labels is None:
            print("警告：未找到标签！")
            print("输入内容：", inputs)
            raise ValueError("未找到标签，请检查数据集是否包含标签。")
        outputs = model(**inputs)
        logits = outputs.get("logits")
        loss_fct = torch.nn.CrossEntropyLoss()
        loss = loss_fct(logits.view(-1, self.model.config.num_labels), labels.view(-1))
        return (loss, outputs) if return_outputs else loss


# 训练模型
# 此函数用于训练模型，接受一个经过划分的数据集作为输入
# 划分出训练集和验证集，设置训练参数，使用自定义的 Trainer 进行训练
def train_model(dataset):
    train_dataset = dataset["train"]
    val_dataset = dataset["test"]

    print("训练集键名:", train_dataset.column_names)  # 调试信息
    print("验证集键名:", val_dataset.column_names)  # 调试信息

    training_args = TrainingArguments(
        output_dir='./results',  # 训练结果输出目录
        num_train_epochs=3,  # 训练轮数
        per_device_train_batch_size=16,  # 每个设备的训练批次大小
        per_device_eval_batch_size=64,  # 每个设备的评估批次大小
        warmup_steps=500,  # 热身步骤数
        weight_decay=0.01,  # 权重衰减率
        logging_dir='./logs',  # 日志输出目录
        logging_steps=10  # 日志记录间隔步数
    )

    trainer = CustomTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=val_dataset,
        tokenizer=tokenizer
    )

    trainer.train()
    return model


# 保存模型
# 将训练好的模型保存到指定路径
def save_model(model, model_path):
    model.save_pretrained(model_path)


if __name__ == "__main__":
    # name替换为想要获取的聊天记录的人的备注
    name = "不可一日无笑"
    path = os.getcwd()
    file_path = os.path.join(path, f"{name}所有聊天记录.txt")

    if not os.path.exists(file_path):
        print(f"聊天记录文件 {file_path} 不存在，请先运行主程序获取聊天记录。")
    else:
        # 读取聊天数据
        chat_data = read_chat_data(file_path)
        # 进行情感分析
        sentiment_data = add_sentiment_analysis(chat_data)
        # 创建数据集
        dataset = create_dataset(chat_data, sentiment_data)

        # 划分训练集和验证集
        dataset = dataset.train_test_split(test_size=0.2)

        # 对数据集进行预处理
        tokenized_dataset = dataset.map(preprocess_function, batched=True)
        # 确保 labels 列被保留
        tokenized_dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'labels'])

        # 训练模型
        trained_model = train_model(tokenized_dataset)

        # 保存模型的路径
        model_path = os.path.join(path, f"{name}_finetuned_model")
        # 保存训练好的模型
        save_model(trained_model, model_path)

        print("模型训练并保存成功。")
